import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchensemble import VotingClassifier
from pointcept.datasets import ScannetDataset
from pointcept.models import PointNet
from pointcept.engines.train import Trainer
from sklearn.model_selection import train_test_split
import numpy as np

# 设定随机种子
np.random.seed(0)
torch.manual_seed(0)

# 数据预处理
# 假设你已经有了一个 ScannetDataset 实例
dataset = ScannetDataset()

# 划分数据集
train_set, val_set = train_test_split(dataset, test_size=0.2, random_state=0)

# 创建数据加载器
batch_size = 128
train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_set, batch_size=batch_size, shuffle=False)

# 定义基分类器
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.pointnet = PointNet(num_classes=20)

    def forward(self, x):
        return self.pointnet(x)

# 初始化集成模型
ensemble = VotingClassifier(
    estimator=Net,
    n_estimators=3,  # 三个基分类器
    cuda=True,
    random_state=0,
)

# 定义训练器
trainer = Trainer(
    ensemble,
    criterion=nn.CrossEntropyLoss(),
    optimizer=lambda model: optim.Adam(model.parameters(), lr=1e-3),
    device='cuda' if torch.cuda.is_available() else 'cpu',
    num_epochs=100,
    logging_period=100,
    eval_dataloader=val_loader,
)

# 训练集成模型
trainer.train(train_loader)

# 验证集成模型
ensemble.eval()
correct = 0
total = 0
with torch.no_grad():
    for data, target in val_loader:
        outputs = ensemble.predict(data)
        _, predicted = torch.max(outputs.data, 1)
        total += target.size(0)
        correct += (predicted == target).sum().item()

print('Accuracy of the network on the validation set: %d %%' % (100 * correct / total))